package com.su02.multi.examples.mlp;

import com.su02.multi.chainrule.Constant;
import com.su02.multi.chainrule.Function;
import com.su02.multi.util.ListUtil;

import org.apache.commons.collections4.CollectionUtils;

import java.util.ArrayList;
import java.util.BitSet;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.PrimitiveIterator;

public final class MultilayerPerceptron extends Function {
    private List<Layer4MLP> layers;

    private MultilayerPerceptron() {}

    public static MultilayerPerceptron makeLayers(List<Integer> structure, List<Double> vector, Iterator<Integer> idGenerator) {
        if (CollectionUtils.isEmpty(structure)) {
            throw new IllegalArgumentException();
        }
        MultilayerPerceptron ans = new MultilayerPerceptron();
        if (!Objects.equals(ListUtil.last(structure), 1)) {
            throw new IllegalArgumentException();
        }
        ans.layers = new ArrayList<>(CollectionUtils.size(structure));

        Layer4MLP firstLayer = Layer4MLP.fromVector(vector, ListUtil.getOrDefault(structure, 1, 1), idGenerator);
        ans.layers.add(firstLayer);

        Layer4MLP lastLayer = firstLayer;
        for (int i = 1 ; i < CollectionUtils.size(structure) ; i ++ ) {
            Layer4MLP layer = Layer4MLP.fromFunctions(lastLayer.out(), structure.get(i), idGenerator);
            ans.layers.add(layer);
            lastLayer = layer;
        }

        BitSet bitSet = new BitSet();

        for (Layer4MLP layer : ans.layers) {
            for (Neuron4MLP f : layer.out()) {
                bitSet.or(f.getBs());
            }
        }

        ans.setBs(bitSet);

        return ans;
    }

    public static Function MSELossFunction(List<Integer> structure, TrainMatrix4MLP matrix) {
        MultilayerPerceptron mlp = makeLayers(structure, matrix.getEntries().get(0).x(), GlobalIdHolder.iterator());
        int nRows = matrix.shape().getLeft();
        List<Integer> lst = new ArrayList<>();
        mlp.getBs().stream().forEach(lst::add);
        Function f = new Constant(0.);
        for (TrainEntry4MLP entry : matrix) {
            f = f .add(MSELossFunctionEntry(structure, entry, lst.iterator()).mul(1. / (2. * nRows)));
        }
        return f;
    }

    private static Function MSELossFunctionEntry(List<Integer> structure, TrainEntry4MLP entry, Iterator<Integer> idGenerator) {
        MultilayerPerceptron mlp = makeLayers(structure, entry.x(), idGenerator);
        return mlp.sub(new Constant(entry.y())).pow(2);
    }

    @Override
    public Function doBackward(int id) {
        return ListUtil.exactOne(ListUtil.last(layers).out()).backward(id);
    }

    @Override
    public double forward(Map<Integer, Double> vector) {
        return ListUtil.exactOne(ListUtil.last(layers).out(vector));
    }
}